from pocketflow import Node
from utils import search_chinanews, call_zhipuai
from typing import Dict, Any, Optional, Tuple, List, Union
import logging
import random

# 导入配置和日志模块（如果尚未创建，则注释掉这些导入）
try:
    from config import get_config
    from logger import AgentLogger
    HAS_CONFIG = True
except ImportError:
    HAS_CONFIG = False
    logging.warning("配置或日志模块导入失败，使用默认配置")

class WebSearchNode(Node):
    """
    网络搜索节点，用于从中新网获取信息
    
    处理流程：
    1. 准备阶段：获取搜索关键词
    2. 执行阶段：执行搜索并格式化结果
    3. 后处理阶段：将结果保存到共享存储并转向决策节点
    """
    def __init__(self, max_retries: int = 2, wait: int = 1):
        """
        初始化搜索节点
        
        参数:
            max_retries: 最大重试次数
            wait: 重试等待时间（秒）
        """
        super().__init__(max_retries=max_retries, wait=wait)
        self._config = get_config()["search"] if HAS_CONFIG else {
            "max_results": 10,
            "timeout": 30
        }
    
    def prep(self, shared: Dict[str, Any]) -> str:
        """
        准备搜索关键词
        
        参数:
            shared: 共享存储字典
            
        返回:
            搜索关键词
        """
        # 获取或初始化日志记录器
        self.logger = shared.get("logger")
        if not self.logger and HAS_CONFIG:
            from config import get_config
            from logger import AgentLogger
            shared["logger"] = AgentLogger(get_config()["log"])
            self.logger = shared["logger"]
            
        # 获取查询
        query = shared.get("search_query")
        if not query:
            query = input("请输入搜索关键词: ")
        
        if self.logger:
            self.logger.info(f"准备搜索: {query}")
        else:
            print(f"🔍 正在搜索: {query}")
            
        return query
    
    def exec(self, prep_res: str) -> str:
        """
        执行搜索并格式化结果
        
        参数:
            prep_res: 搜索关键词
            
        返回:
            格式化的搜索结果
        """
        if not prep_res:
            error_msg = "搜索关键词为空"
            if hasattr(self, "logger") and self.logger:
                self.logger.error(error_msg)
            return error_msg
            
        # 获取最大结果数
        max_results = self._config.get("max_results", 10)
        
        # 执行搜索
        results = search_chinanews(prep_res, num_results=max_results)
        
        # 记录搜索结果
        if hasattr(self, "logger") and self.logger:
            try:
                # 将Pydantic模型转换为字典以便记录
                results_dicts = [result.dict() for result in results]
                self.logger.log_search(prep_res, results_dicts)
            except Exception as e:
                self.logger.error(f"记录搜索结果时出错: {e}")
            
        if not results:
            # 使用标准格式的无结果消息，确保post方法能够识别
            no_results_msg = f"未找到与'{prep_res}'相关的信息。建议尝试其他关键词搜索。"
            if hasattr(self, "logger") and self.logger:
                self.logger.warning("搜索无结果")
            return no_results_msg
            
        # 拼接结果为文本，改进格式提高可读性
        context = "\n\n".join([
            f"【{i+1}】 {result.title}\n链接: {result.url}" 
            for i, result in enumerate(results)
        ])
        
        return context

    def post(self, shared: Dict[str, Any], prep_res: Optional[str], exec_res: Optional[str]) -> str:
        """
        保存结果到共享存储并确定下一步
        
        参数:
            shared: 共享存储字典
            prep_res: 搜索关键词
            exec_res: 搜索结果
            
        返回:
            下一步操作的名称
        """
        if prep_res is None or exec_res is None:
            error_msg = "搜索失败，无法获取结果"
            if hasattr(self, "logger") and self.logger:
                self.logger.error(error_msg)
            else:
                print(f"❌ {error_msg}")
            return None
            
        # 保存上下文并记录
        shared["context"] = exec_res
        shared["last_search_query"] = prep_res
        
        # 检查是否为空搜索结果
        if "未找到" in exec_res and "相关信息" in exec_res:
            warning_msg = f"搜索完成，但未找到相关信息: {prep_res}"
            if hasattr(self, "logger") and self.logger:
                self.logger.warning(warning_msg)
            else:
                print(f"⚠️ {warning_msg}")
        else:
            success_msg = "搜索完成，找到相关信息"
            if hasattr(self, "logger") and self.logger:
                self.logger.info(success_msg)
            else:
                print(f"✅ {success_msg}")

        # 返回下一步操作，必须与Flow中连接的动作名称一致
        return "decide"

class DecideNode(Node):
    """
    决策节点，用于决定下一步操作
    
    处理流程:
    1. 准备阶段：获取问题和当前检索上下文
    2. 执行阶段：调用LLM决定下一步动作
    3. 后处理阶段：根据决策结果更新共享存储并决定流程方向
    """
    def __init__(self, max_retries: int = 2, wait: int = 1):
        """
        初始化决策节点
        
        参数:
            max_retries: 最大重试次数
            wait: 重试等待时间（秒）
        """
        super().__init__(max_retries=max_retries, wait=wait)
        self._config = get_config()["llm"] if HAS_CONFIG else {
            "temperature": 0.7,
            "max_tokens": 1024
        }
        self._flow_config = get_config()["flow"] if HAS_CONFIG else {
            "max_search_iterations": 5
        }
        
    def prep(self, shared: Dict[str, Any]) -> Tuple[str, str]:
        """
        准备question和context
        
        参数:
            shared: 共享存储字典
            
        返回:
            问题和检索上下文的元组
        """
        # 获取记录器
        self.logger = shared.get("logger")
        
        # 检查搜索迭代次数
        if "search_iterations" not in shared:
            shared["search_iterations"] = 0
            
        # 从共享存储中获取问题，如果没有则提示输入
        question = shared.get("question")
        if not question:
            question = input("请输入问题: ")
            shared["question"] = question
            
        context = shared.get("context", "没有检索信息")
        
        if self.logger:
            self.logger.debug(f"准备决策，问题: {question[:30]}..., 上下文长度: {len(context)}")
            
        return question, context
    
    def exec(self, prep_res: Tuple[str, str]) -> Dict[str, Any]:
        """
        调用LLM决定下一步操作
        
        参数:
            prep_res: 包含问题和上下文的元组
            
        返回:
            包含决策信息的字典
        """
        question, context = prep_res
        
        # 检查搜索迭代次数是否超过限制
        if hasattr(self, "shared") and self.shared:
            search_iterations = self.shared.get("search_iterations", 0)
            max_iterations = self._flow_config.get("max_search_iterations", 5)
            
            # 获取连续无结果搜索计数
            no_results_count = self.shared.get("no_results_count", 0)
            
            if search_iterations >= max_iterations:
                if self.logger:
                    self.logger.warning(f"达到最大搜索次数 {max_iterations}，强制回答")
                else:
                    print(f"⚠️ 达到最大搜索次数 {max_iterations}，强制回答")
                    
                return {
                    "action": "answer",
                    "reason": f"已达到最大搜索次数限制 ({max_iterations})",
                    "answer": "无法找到完整答案，但以下是基于现有信息的尝试回答: " + context
                }
            
            # 检查当前上下文是否为空搜索结果
            is_no_results = False
            # 使用多种标识模式检查是否为空结果
            if any(pattern in context for pattern in ["未找到相关信息", "未找到与", "没有找到", "无相关结果"]):
                is_no_results = True
                # 增加无结果计数
                self.shared["no_results_count"] = no_results_count + 1
                
                # 如果连续3次无结果，直接转向回答
                if no_results_count >= 2:  # 已经有2次，这是第3次
                    if self.logger:
                        self.logger.warning(f"连续{no_results_count + 1}次搜索无结果，转向回答")
                    else:
                        print(f"⚠️ 连续{no_results_count + 1}次搜索无结果，转向回答")
                        
                    return {
                        "action": "answer",
                        "reason": f"连续多次搜索无相关结果",
                        "answer": "经过多次搜索，未能找到与您问题直接相关的信息。以下是基于现有知识的回答: "
                    }
            else:
                # 重置无结果计数
                self.shared["no_results_count"] = 0
                
            if is_no_results and search_iterations > 0 and no_results_count < 3:
                # 如果已经搜索过，并且没有找到结果，尝试使用不同关键词
                if self.logger:
                    self.logger.info(f"上次搜索未找到结果(第{no_results_count}次)，尝试使用不同关键词")
                else:
                    print(f"🔄 上次搜索未找到结果(第{no_results_count}次)，尝试使用不同关键词")
                
                # 获取上一次使用的搜索词
                last_query = self.shared.get("last_search_query", "")
                
                # 直接返回搜索决策，不再调用LLM
                return {
                    "action": "search",
                    "reason": "上次搜索未找到结果，尝试使用更广泛的关键词",
                    "search_query": self._generate_alternative_query(question, last_query, no_results_count)
                }
        
        if self.logger:
            self.logger.info("Agent正在思考如何决策...")
        else:
            print("🤔 Agent正在思考如何决策...")
        
        # 在提示中添加关于无结果情况的处理建议及搜索工具特性说明
        prompt = f"""
### 上下文
你是一个信息检索专家，现在需要根据问题和已有的检索信息，决定下一步做什么。
问题：{question}
已有的检索信息：{context}

### 搜索工具特性说明
1. 当前搜索工具是中国新闻网(chinanews)的搜索引擎，主要针对新闻报道进行检索
2. 该工具适合查找新闻事件、政策公告、企业动态等新闻类内容
3. 不适合查找技术文档、学术论文、专业研究报告等非新闻类内容
4. 搜索结果仅限于该新闻网站收录的内容，不是全网搜索

### 搜索策略建议
1. 使用具体新闻事件相关词汇，如"发布"、"宣布"、"报道"等
2. 添加具体公司、机构名称，如"中国航天科技集团"、"中国卫星通信公司"等
3. 添加新闻热点词汇，如"发射"、"合作"、"计划"等
4. 避免使用过于专业或学术的词汇
5. 如果多次搜索无结果，建议转向更通用的新闻话题或直接回答问题

### 行为选项
[1] search
    description: 继续搜索相关信息
    parameters:
        query: 新的搜索关键词

[2] answer
    description: 根据已有的检索信息，给出问题的答案
    parameters:
        answer: 问题的答案


### 决策
请选择一个行为，并使用 yaml 格式返回。
示例：

```yaml
thinking: |
    <你一步一步推理的过程>
action: search or answer
reason: <你选择该行为的理由>
answer: <如果选择的行动是answer，则返回问题的答案>
search_query: <如果选择的行动是search，则返回新的搜索关键词>
```

注意: 请确保:
1. 使用适当的缩进（4个空格）用于所有多行字段
2. 使用 | 字符表示多行文本字段
3. 保持单行字段没有 | 字符

如果检索信息提示"未找到相关信息"，请选择search动作并提供一个新的、更合适的搜索关键词。
可以尝试使用更通用的词语、缩短关键词、或者使用近义词来扩大搜索范围。
如果连续多次搜索无结果，请考虑转向answer动作，基于你的知识给出回答。
"""

        # 调用LLM
        response = call_zhipuai(prompt)

        # 解析yaml
        import yaml
        try:
            # 尝试提取yaml部分
            if "```yaml" in response and "```" in response.split("```yaml")[1]:
                yaml_str = response.split("```yaml")[1].split("```")[0].strip()
            elif "```" in response and "```" in response.split("```")[1]:
                # 可能没有指定yaml，但使用了代码块
                yaml_str = response.split("```")[1].split("```")[0].strip()
            else:
                # 没有代码块，尝试直接解析整个响应
                yaml_str = response.strip()
                
            # 尝试加载YAML
            try:
                decision = yaml.safe_load(yaml_str)
                
                # 确保决策包含必要的字段
                if not isinstance(decision, dict):
                    raise ValueError("决策必须是字典格式")
                    
                if "action" not in decision:
                    # 尝试从文本中推断动作
                    if "search" in yaml_str.lower() and "query" in yaml_str.lower():
                        decision["action"] = "search"
                        if "search_query" not in decision:
                            # 尝试提取查询
                            search_lines = [line for line in yaml_str.split("\n") if "query" in line.lower()]
                            if search_lines:
                                query_value = search_lines[0].split(":", 1)[1].strip() if ":" in search_lines[0] else question
                                decision["search_query"] = query_value
                            else:
                                decision["search_query"] = question
                    else:
                        decision["action"] = "answer"
                        if "answer" not in decision:
                            decision["answer"] = "基于现有信息无法给出明确答案。" + response
                
            except yaml.YAMLError:
                # 如果YAML解析失败，尝试手动提取决策
                print("YAML解析失败，尝试手动提取决策...")
                decision = {"thinking": response}
                
                if "search" in response.lower() and ("搜索" in response or "查询" in response):
                    decision["action"] = "search"
                    # 尝试提取搜索查询
                    search_lines = [line for line in response.split("\n") if "搜索" in line or "查询" in line or "search_query" in line or "query" in line]
                    if search_lines:
                        for line in search_lines:
                            if ":" in line:
                                decision["search_query"] = line.split(":", 1)[1].strip()
                                break
                        if "search_query" not in decision:
                            decision["search_query"] = question
                    else:
                        decision["search_query"] = question
                else:
                    decision["action"] = "answer"
                    decision["answer"] = "根据现有信息，" + response
                
                decision["reason"] = "从文本中提取的决策"
            
            # 记录决策
            if self.logger:
                self.logger.log_decision(decision)
                
            return decision
                
        except Exception as e:
            error_msg = f"解析决策失败: {e}"
            if self.logger:
                self.logger.error(error_msg)
            else:
                print(f"⚠️ {error_msg}")
                
            # 提供默认返回，避免流程中断
            return {
                "action": "search",
                "reason": "默认决策，因为解析失败",
                "search_query": question
            }

    def post(self, shared: Dict[str, Any], prep_res: Tuple[str, str], exec_res: Dict[str, Any]) -> str:
        """
        保存决策并确定流程的下一步
        
        参数:
            shared: 共享存储字典
            prep_res: 包含问题和上下文的元组
            exec_res: 决策结果
            
        返回:
            下一步操作的名称
        """
        # 保存共享内容以便在exec中访问
        self.shared = shared
        
        if not exec_res:
            error_msg = "决策失败"
            if self.logger:
                self.logger.error(error_msg)
            else:
                print(f"❌ {error_msg}")
            return None
            
        # 显示思考过程
        if "thinking" in exec_res and get_config()["ui"]["show_thinking"] if HAS_CONFIG else True:
            print(f"\n🧠 思考过程:\n{exec_res['thinking']}\n")
            
        reason = exec_res.get("reason", "未提供")
        if self.logger:
            self.logger.debug(f"决策理由: {reason}")
        else:
            print(f"🔍 决策理由: {reason}")
            
        # 处理决策结果
        action = exec_res.get("action", "search")
        
        if action == "search":
            # 递增搜索次数
            shared["search_iterations"] = shared.get("search_iterations", 0) + 1
            
            search_query = exec_res.get("search_query")
            if not search_query:
                search_query = prep_res[0]  # 如果没有提供搜索词，使用原问题
                
            shared["search_query"] = search_query
            
            if self.logger:
                self.logger.info(f"Agent决定搜索: {search_query}")
            else:
                print(f"🔍 Agent决定搜索: {search_query}")
                
        elif action == "answer":
            answer = exec_res.get("answer", "未能提供答案")
            shared["answer_content"] = answer  # 存储答案内容
            
            if self.logger:
                self.logger.info("Agent决定回答问题")
            else:
                print(f"💡 Agent决定回答问题")
                
        else:
            error_msg = f"未知操作: {action}，默认执行搜索"
            if self.logger:
                self.logger.warning(error_msg)
            else:
                print(f"⚠️ {error_msg}")
                
            action = "search"
            shared["search_query"] = prep_res[0]
            
        # 返回动作以确定流程的下一步
        return action

    def _generate_alternative_query(self, question: str, last_query: str, attempt: int = 1) -> str:
        """
        生成替代搜索查询，当上一次搜索失败时使用
        
        参数:
            question: 原始问题
            last_query: 上一次使用的搜索查询
            attempt: 尝试次数，影响策略选择
            
        返回:
            新的搜索查询
        """
        # 从原问题中提取关键词
        keywords = question.replace("？", "").replace("?", "").split()
        
        # 新闻相关词汇，用于增强查询的新闻性
        news_verbs = ["报道", "宣布", "发布", "公开", "透露", "披露", "启动"]
        news_nouns = ["消息", "进展", "动态", "规划", "新闻", "计划", "展望"]
        companies = ["中国航天科技集团", "中国卫星通信公司", "中国电子科技集团", "华为", "中国移动", "中国电信"]
        
        # 按照尝试次数选择不同的策略
        if attempt == 1:
            # 第一次尝试：添加新闻动词
            news_verb = random.choice(news_verbs)
            
            # 如果问题中包含"发展现状"，将其替换为更具新闻性的词汇
            if "发展现状" in question or "发展状况" in question:
                query = question.replace("发展现状", "最新进展").replace("发展状况", "最新动态")
                return query + " " + news_verb
            else:
                # 提取问题核心，加上新闻词汇
                return " ".join(keywords[:3]) + " " + news_verb
        
        elif attempt == 2:
            # 第二次尝试：使用实体名称
            # 检查是否包含特定领域
            if "卫星" in question or "航天" in question:
                relevant_companies = [c for c in companies if "航天" in c or "卫星" in c]
                if relevant_companies:
                    return relevant_companies[0] + " " + random.choice(news_nouns)
            
            # 如果没有特定领域匹配，使用更通用的新闻查询
            return "中国" + random.choice(["卫星网络", "通信卫星", "卫星互联网"]) + " " + random.choice(news_nouns)
                    
        else:
            # 第三次或更多尝试：使用极简、通用的新闻查询
            generic_queries = [
                "卫星通信 最新动态",
                "中国航天 卫星计划",
                "卫星网络 应用",
                "航天科技 新闻",
                "通信卫星 发射"
            ]
            
            # 避免重复之前使用过的查询
            for query in generic_queries:
                if query != last_query:
                    return query
                
        # 最后的回退方案
        return "卫星通信 新闻"

class AnswerQuestionNode(Node):
    """
    回答问题节点，用于生成最终答案
    
    处理流程:
    1. 准备阶段：获取问题和上下文
    2. 执行阶段：调用LLM生成答案
    3. 后处理阶段：保存答案并完成流程
    """
    def __init__(self, max_retries: int = 2, wait: int = 1):
        """
        初始化回答节点
        
        参数:
            max_retries: 最大重试次数
            wait: 重试等待时间（秒）
        """
        super().__init__(max_retries=max_retries, wait=wait)
        self._config = get_config()["llm"] if HAS_CONFIG else {
            "temperature": 0.7,
            "max_tokens": 1024
        }
    
    def prep(self, shared: Dict[str, Any]) -> Tuple[str, str]:
        """
        获取问题和上下文以回答问题
        
        参数:
            shared: 共享存储字典
            
        返回:
            问题和上下文的元组
        """
        # 获取记录器
        self.logger = shared.get("logger")
        
        question = shared.get("question", "未指定问题")
        
        # 如果决策节点已经生成了答案，直接使用
        if "answer_content" in shared:
            context = shared["answer_content"]
        else:
            context = shared.get("context", "")
            
        if self.logger:
            self.logger.debug(f"准备回答，问题: {question[:30]}..., 上下文长度: {len(context)}")
            
        return question, context
        
    def exec(self, prep_res: Tuple[str, str]) -> str:
        """
        调用LLM生成最终答案
        
        参数:
            prep_res: 问题和上下文的元组
            
        返回:
            生成的答案
        """
        question, context = prep_res
        
        if self.logger:
            self.logger.info("正在生成最终答案...")
        else:
            print(f"✍️ 正在生成最终答案...")
        
        # 如果决策节点已经生成了答案，直接返回
        if "answer_content" in context:
            return context
        
        # 检查上下文是否为空搜索结果
        is_no_results = False
        if any(pattern in context for pattern in ["未找到相关信息", "未找到与", "没有找到", "无相关结果"]):
            is_no_results = True
            if self.logger:
                self.logger.warning("生成答案时检测到搜索无结果")
            else:
                print("⚠️ 生成答案时检测到搜索无结果")
                
        # 根据有无结果选择不同的提示模板
        if is_no_results:
            prompt = f"""
### 上下文
尝试回答以下问题，但在中国新闻网的搜索中未找到相关的新闻报道。
问题：{question}

## 你的回答：
请根据你的知识库提供关于该问题的信息。在回答中：

1. 明确说明这不是基于中国新闻网的检索结果，而是基于你的知识
2. 提供关于主题的主要发展方向、重要项目和成就
3. 提及主要的参与公司、机构和关键人物
4. 讨论相关技术的主要应用领域
5. 如果有相关限制，请说明你的信息可能不包含最新进展

请以客观、清晰的方式组织信息，避免过于技术性的细节。
"""
        else:
            # 创建提示以生成答案
            prompt = f"""
### 上下文
根据以下信息，回答问题。
问题：{question}
检索信息：{context}

## 你的回答：
请提供详细、准确、有条理的回答。如果检索信息不足，请明确指出并说明需要更多信息。
"""
        
        # 调用LLM生成答案
        answer = call_zhipuai(prompt)
        return answer
    
    def post(self, shared: Dict[str, Any], prep_res: Tuple[str, str], exec_res: str) -> str:
        """
        保存最终答案并完成流程
        
        参数:
            shared: 共享存储字典
            prep_res: 问题和上下文的元组
            exec_res: 生成的答案
            
        返回:
            流程结束标记
        """
        if not exec_res:
            error_msg = "答案生成失败"
            if self.logger:
                self.logger.error(error_msg)
            else:
                print(f"❌ {error_msg}")
            return None
            
        # 记录回答
        question = prep_res[0]
        if self.logger:
            self.logger.log_answer(question, exec_res)
            
        # 在共享存储中保存答案
        shared["answer"] = exec_res
        
        # 打印答案
        print(f"\n📝 最终答案:\n{exec_res}\n")
        
        if self.logger:
            self.logger.info("回答生成成功")
            # 保存历史记录
            self.logger.save()
        else:
            print(f"✅ 回答生成成功")
        
        # 流程结束
        return "done"
        

